#pragma once
#include <chrono>
#include <iostream>
#include <string>
#include <torch/torch.h>
#include <vector>

// 打印张量信息
inline void prtts1(const torch::Tensor &tensor, const std::string &name) {
  if (!name.empty()) {
    std::cout << name << ": ";
  }
  std::cout << "shape=" << tensor.sizes() << ", dtype=" << tensor.dtype()
            << ", device=" << tensor.device()
            << ", requires_grad=" << (tensor.requires_grad() ? "true" : "false")
            << std::endl;
  std::cout << "values:\n" << tensor << std::endl;
}

// 打印张量信息
inline void prtts2(const torch::Tensor &tensor,
                   const std::string &name = "Tensor") {
  std::cout << "==== " << name << " ====\n";
  std::cout << tensor << "\n";
}

// 常用工具函数
namespace pytorch_study {

// 创建示例张量
inline torch::Tensor create_sample_tensor(int64_t rows, int64_t cols) {
  return torch::arange(rows * cols).reshape({rows, cols}).to(torch::kFloat32);
}

// 计时器类
class Timer {
public:
  Timer(const std::string &name = "") : name_(name) {
    start_ = std::chrono::high_resolution_clock::now();
  }

  ~Timer() {
    auto end = std::chrono::high_resolution_clock::now();
    auto duration =
        std::chrono::duration_cast<std::chrono::milliseconds>(end - start_);
    std::cout << name_ << " took " << duration.count() << " ms" << std::endl;
  }

private:
  std::string name_;
  std::chrono::time_point<std::chrono::high_resolution_clock> start_;
};

} // namespace pytorch_study
